import python_ai.ML.logic_regression.multiple_classification.the_data as data
from python_ai.ML.logic_regression.xlib_logic_regression import *
from python_ai.ML.lin_regression.xlib import *
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

np.random.seed(1)
plt.figure(figsize=[12, 8])
spr = 1
spc = 2
plt1 = plt.subplot(spr, spc, 1)
plt2 = plt.subplot(spr, spc, 2)

m = len(data.x1x2y)
for k, xpara in enumerate(data.xparams):
    cls_id = k + 1

    x1x2 = data.x1x2y[:, 0:-1]
    x1x2y = np.c_[x1x2, np.zeros([m, 1])]
    cls_selector = (data.x1x2y[:, -1] == cls_id).ravel()
    m_selector = np.sum(cls_selector, axis=0)
    cls_unselector = np.invert(cls_selector)
    m_unselector = m - m_selector
    x1x2y[cls_selector, -1] = np.ones(m_selector)
    y = x1x2y[:, -1]
    x1_unscaled = x1 = x1x2y[:, 0].reshape(m, 1)
    x2_unscaled = x2 = x1x2y[:, 1].reshape(m, 1)

    # plot original data
    plt_x = x1[cls_selector].ravel()
    plt_y = x2[cls_selector].ravel()
    x1x2_ori = np.c_[x1, x2]
    plt1.scatter(plt_x, plt_y, s=4, label='cls ' + str(cls_id))

    # scale data
    std_scaler = StandardScaler()
    std_scaler.fit(x1x2_ori)
    mu = std_scaler.mean_
    sigma = std_scaler.scale_
    x1x2_scaled = std_scaler.transform(x1x2_ori)
    x1_scaled = x1x2_scaled[:, 0].reshape(m, 1)
    x2_scaled = x1x2_scaled[:, 1].reshape(m, 1)

    # plot scaled data
    plt_x_scaled = x1_scaled[cls_selector].ravel()
    plt_y_scaled = x2_scaled[cls_selector].ravel()
    plt2.scatter(plt_x_scaled, plt_y_scaled, s=4, label='scaled cls ' + str(cls_id))

    # logic regression
    model = LogisticRegression(solver='liblinear')
    model.fit(x1x2_scaled, y)
    theta0 = model.intercept_.reshape(1, 1)
    theta1_n = model.coef_.reshape(2, 1)
    theta = np.r_[theta0, theta1_n]
    print(f'THETA({cls_id}) = {theta}')

    # the decision border
    plt_theta_scaled = transform_theta_of_logic_regression(theta)
    plt_line_x_scaled = np.array([x1_scaled.min(axis=0), x1_scaled.max(axis=0)]).ravel()
    mm = len(plt_line_x_scaled)
    plt_line_y_scaled = np.c_[np.ones([mm, 1]), plt_line_x_scaled.reshape(mm, 1)].dot(plt_theta_scaled).ravel()
    plt2.plot(plt_line_x_scaled, plt_line_y_scaled, label='scaled board of cls ' + str(cls_id))

    plt_theta_unscaled = scale_theta_back(x1_unscaled, x2_unscaled, plt_theta_scaled)
    plt_line_x_unscaled = np.array([x1_unscaled.min(axis=0), x1_unscaled.max(axis=0)]).ravel()
    mm = len(plt_line_x_unscaled)
    plt_line_y_unscaled = np.c_[np.ones([mm, 1]), plt_line_x_unscaled.reshape(mm, 1)].dot(plt_theta_unscaled).ravel()
    plt1.plot(plt_line_x_unscaled, plt_line_y_unscaled, label='unscaled board of cls ' + str(cls_id))

plt1.grid()
plt1.legend()
plt2.grid()
plt2.legend()
plt.show()
